import os
import random
from PIL import Image, ImageOps

# ============================
# 基础配置
# ============================
# input_dir = "E:\output_image"             # 原始图片文件夹
input_dir = "image"
output_dir = "dataset"           # 输出数据集目录
train_txt = os.path.join(output_dir, "train.txt")
val_txt = os.path.join(output_dir, "val.txt")

# 训练集与验证集划分比例
train_ratio = 0.8

# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)

# ============================
# 辅助函数
# ============================
def save_rotated_images(img_path, index):
    """
    保存原图与旋转图像
    返回 [(filename, label), ...]
    """
    print(f"正在处理第{index}张 {img_path}")
    img = Image.open(img_path)
    img = ImageOps.exif_transpose(img)
    img = img.convert("RGB")  # 确保统一为RGB
    base_name = f"{index:05d}"

    samples = []

    # 0次旋转（正面）
    name_0 = f"0_cls_{base_name}.jpg"
    img.save(os.path.join(output_dir, name_0))
    samples.append((name_0, 0))

    # 顺时针旋转 90° (标签3)
    name_3 = f"3_cls_{base_name}.jpg"
    img_3 = img.rotate(-270, expand=True)
    img_3.save(os.path.join(output_dir, name_3))
    samples.append((name_3, 3))

    # 顺时针旋转 180° (标签2)
    name_2 = f"2_cls_{base_name}.jpg"
    img_2 = img.rotate(-180, expand=True)
    img_2.save(os.path.join(output_dir, name_2))
    samples.append((name_2, 2))

    # 顺时针旋转 270° (标签1)
    name_1 = f"1_cls_{base_name}.jpg"
    img_1 = img.rotate(-90, expand=True)
    img_1.save(os.path.join(output_dir, name_1))
    samples.append((name_1, 1))

    return samples


# ============================
# 主逻辑
# ============================
def main():
    image_files = [f for f in os.listdir(input_dir)
                   if f.lower().endswith((".jpg", ".jpeg", ".png", ".bmp"))]

    all_samples = []
    for i, filename in enumerate(image_files):
        full_path = os.path.join(input_dir, filename)
        samples = save_rotated_images(full_path, i + 1)
        all_samples.extend(samples)

    # 打乱样本
    random.shuffle(all_samples)

    # 划分训练集与验证集
    split_idx = int(len(all_samples) * train_ratio)
    train_samples = all_samples[:split_idx]
    val_samples = all_samples[split_idx:]

    # 写入txt
    with open(train_txt, "w", encoding="utf-8") as f:
        for name, label in train_samples:
            f.write(f"{os.path.basename(output_dir)}/{name} {label}\n")

    with open(val_txt, "w", encoding="utf-8") as f:
        for name, label in val_samples:
            f.write(f"{os.path.basename(output_dir)}/{name} {label}\n")

    print(f"✅ 处理完成！共生成 {len(all_samples)} 张样本：")
    print(f"  训练集: {len(train_samples)} 张")
    print(f"  验证集: {len(val_samples)} 张")
    print(f"  输出目录: {os.path.abspath(output_dir)}")


if __name__ == "__main__":
    main()
